/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.
 * Description: 节点/设备间的链接派生类头文件
 * Author: lzh
 * Create: 2024-07-27
 */

#ifndef TRANSPORT_DEVICE_IBVERBS_PUB_H
#define TRANSPORT_DEVICE_IBVERBS_PUB_H

#include <functional>
#include <atomic>
#include "transport_ibverbs_pub.h"

namespace hccl {

using WrInfo = struct TagWrInfo {
    struct wr_info wrData;
    u64 type;
    u64 wrDataAddr;
    TagWrInfo() : type(static_cast<u64>(WqeType::WQE_TYPE_DATA)), wrDataAddr(0)
    {
        wrData = {0};
    }
};

constexpr u32 IBV_SGLIST_LEN_MAX = 2147483648;

class TransportDeviceIbverbs : public TransportIbverbs {
public:
    TransportDeviceIbverbs(DispatcherPub *dispatcher,
                           const std::unique_ptr<NotifyPool> &notifyPool,
                           MachinePara &machinePara,
                           std::chrono::milliseconds timeout,
                           const TransportDeviceIbverbsData &transDevIbverbsData);
    ~TransportDeviceIbverbs() override;

    HcclResult Init() override;

    HcclResult TxAsync(UserMemType dstMemType, u64 dstOffset, const void *src, u64 len,
                                  Stream &stream) override;
    HcclResult TxAsync(std::vector<TxMemoryInfo>& txMems, Stream &stream) override;

    HcclResult RxAsync(UserMemType srcMemType, u64 srcOffset, void *dst, u64 len,
                                  Stream &stream) override;
    HcclResult RxAsync(std::vector<RxMemoryInfo>& rxMems, Stream &stream) override;

    HcclResult DataReceivedAck(Stream &stream) override;

    HcclResult TxAck(Stream &stream) override;
    HcclResult RxAck(Stream &stream) override;

    HcclResult TxDataSignal(Stream &stream) override;
    HcclResult RxDataSignal(Stream &stream) override;

    HcclResult TxWaitDone(Stream &stream) override;
    HcclResult TxWithReduce(UserMemType dstMemType, u64 dstOffset, const void *src, u64 len,
                              const HcclDataType datatype, HcclReduceOp redOp, Stream &stream) override;
    HcclResult TxWithReduce(const std::vector<TxMemoryInfo> &txWithReduceMems, const HcclDataType datatype,
        HcclReduceOp redOp, Stream &stream) override;

    HcclResult TxPrepare(Stream &stream) override;
    HcclResult RxPrepare(Stream &stream) override;

    HcclResult TxData(UserMemType dstMemType, u64 dstOffset, const void *src, u64 len,
                                  Stream &stream) override;
    HcclResult RxData(UserMemType srcMemType, u64 srcOffset, void *dst, u64 len,
                                  Stream &stream) override;

    HcclResult TxDone(Stream &stream) override;
    HcclResult RxDone(Stream &stream) override;

    HcclResult WriteAsync(struct Transport::Buffer &remoteBuf, struct Transport::Buffer &localBuf, Stream &stream);

    HcclResult WriteReduceAsync(struct Transport::Buffer &remoteBuf, struct Transport::Buffer &localBuf,
        const HcclDataType datatype, HcclReduceOp redOp, Stream &stream);

    HcclResult PostReady(Stream &stream);
    HcclResult WaitReady(Stream &stream);

    HcclResult PostFin(Stream &stream);
    HcclResult WaitFin(Stream &stream);

    HcclResult PostFinAck(Stream &stream);
    HcclResult WaitFinAck(Stream &stream);

    HcclResult Post(u32 notifyIdx, Stream &stream) override;
    HcclResult Wait(u32 notifyIdx, Stream &stream) override;

    HcclResult AddWrList(void *dstMemPtr, const void *srcMemPtr, u64 srcMemSize, u32 srcKey, u32 dstKey,
        WqeType wqeType, wr_aux_info &aux, std::vector<WrInfo> &wrInfoVec);
    HcclResult GetMemInfo(UserMemType memType, void **dstMemPtr, unsigned int *dstKey, u64 &dstMemSize);
    HcclResult TxPayLoad(UserMemType dstMemType, u64 dstOffset, const void *src, u64 len,
        WqeType wqeType, wr_aux_info &aux, std::vector<WrInfo>& wrInfoVec);
    HcclResult TxSendDataAndNotifyWithSingleQP(std::vector<WrInfo> &wrInfoVec,
        Stream &stream, bool useOneDoorbell = false);
    HcclResult TxSendDataAndNotify(std::vector<WrInfo> &wrInfoVec, Stream &stream, bool useOneDoorbell = false);
    HcclResult TxWrList(std::vector<struct wr_info> &wrlisDataVec, Stream &stream,
        std::vector<struct send_wr_rsp> &opRspVec, u32 multiQpIndex = RDMA_INVALID_QP_INDEX);
    HcclResult SendWrList(
        u32 wrNum, struct wr_info *wrlist, struct send_wr_rsp *opRsp, u32 multiQpIndex = RDMA_INVALID_QP_INDEX);
    HcclResult SendWrlistExt(struct wr_info wr[], struct send_wr_rsp opRsp[], unsigned int sendNum,
        unsigned int *completeNum, u32 multiQpIndex = RDMA_INVALID_QP_INDEX);
    HcclResult TxSendWrlistExt(struct wr_info wrList[], u32 sendNum, struct send_wr_rsp opRsp[],
        unsigned int *completeNum, u32 multiQpIndex = RDMA_INVALID_QP_INDEX);
    HcclResult RdmaSendAsync(struct send_wr &wr, Stream &stream, WqeType wqeType, u64 notifyAddr);
    HcclResult RdmaSendAsync(std::vector<WrInfo> &wqeInfoVec, Stream &stream, bool useOneDoorbell = false,
        u32 multiQpIndex = RDMA_INVALID_QP_INDEX);
    HcclResult GetWrDataAddr(void *dstAddr, WqeType wqeType, u64 &wrDataAddr);
    HcclResult TxSendWqe(void *dstMemPtr, u32 dstKey, const void *srcMemPtr, u32 srcKey,
        u64 srcMemSize, Stream &stream, WqeType wqeType);

    HcclResult ConstructPayLoadWqe(void *dstMemPtr, u32 dstKey, const void *src, u32 srcKey, u64 len,
        WqeType wqeType, wr_aux_info &aux, std::vector<WrInfo> &wrInfoVec, u32 txSendDataTimes);
    HcclResult WriteCommon(const void *remoteAddr, const void *localAddr, u64 length, Stream &stream,
        WqeType wqeType, struct wr_aux_info &aux);
    bool UseMultiQp();
    HcclResult TxSendDataAndNotifyWithMultiQP(
        std::vector<WrInfo> &wqeInfoVec, u32 actualMultiQpNum, Stream &stream, bool useOneDoorbell);
    u32 GetActualQpNum(u32 maxLength);
    HcclResult GetTransportId(u32 &id) override;
private:
    HcclResult CalcDbAddr(u32 dbIndex, u64 &dbAddr);
    
    TransportDeviceIbverbsData transDevIbverbsData_;
    void *notifyValueAddr_ = nullptr;
    MemDetails localInputMem_;
    MemDetails localOutputMem_;
    s64 chipId_;
    static std::atomic<u64> wrIdOffset_;
    u32  multiQpThreshold_{HCCL_MULTI_QP_THRESHOLD_DEFAULT};
};
}  // namespace hccl

#endif /* TRANSPORT_DEVICE_IBVERBS_PUB_H */
